#!/usr/bin/env python
import numpy as np
import sys
import os 
import matplotlib.pyplot as plt
import scienceplots
import argparse

# # generate uav position and orientation
# if os.name == 'nt':  # Windows
#     lib_path = r"c:\Users\ljn\Documents\papers\Dual UWB localization and analysis\simulation"
# elif os.name == 'posix':  # macOS 或 Linux
#     lib_path = "/Users/lvjunning/Nutstore Files/papers/dual UWB localization and analysis/simulation"
# else:
#     raise EnvironmentError("Unsupported Operating System")
# # print(f"Library path: {lib_path}")
# sys.path.append(lib_path)

from range_estimator_4d import heading_filter_tightly_couple_uwb

from rosbag_plot import load_orientation_data, plot_one_dimension_data, load_json_file

# Function to find the closest timestamp
def find_closest_time(reference_time, target_times):
    """
    Find the closest timestamp from target_times to the reference_time.
    """
    return min(target_times, key=lambda t: abs(t - reference_time))

# Synchronize data based on uwb1 timestamps
def synchronize_data(uwb1_data, uwb2_data, uav_data, usv_data):
    """
    Synchronize uwb2, uav, and usv data to match uwb1 timestamps.
    """
    # Extract timestamps
    uwb1_times = [entry['time'] for entry in uwb1_data]
    uwb2_times = [entry['time'] for entry in uwb2_data]
    uav_times = [entry['time'] for entry in uav_data]
    usv_times = [entry['time'] for entry in usv_data]

    synchronized_data = []

    for uwb1_entry in uwb1_data:
        ref_time = uwb1_entry['time']
        
        # Find closest matches
        closest_uwb2 = find_closest_time(ref_time, uwb2_times)
        closest_uav = find_closest_time(ref_time, uav_times)
        closest_usv = find_closest_time(ref_time, usv_times)
        
        # Get corresponding entries
        uwb2_entry = next(entry for entry in uwb2_data if entry['time'] == closest_uwb2)
        uav_entry = next(entry for entry in uav_data if entry['time'] == closest_uav)
        usv_entry = next(entry for entry in usv_data if entry['time'] == closest_usv)
        
        # Combine into a single synchronized entry
        synchronized_data.append({
            'uwb1': uwb1_entry,
            'uwb2': uwb2_entry,
            'uav': uav_entry,
            'usv': usv_entry
        })
    
    return synchronized_data

def save_result(current_dir, result):
    est_file_name = os.path.join(current_dir, "localization_result_{}_anchors.txt".format(6))
    i = 1
    while os.path.exists(est_file_name):
        est_file_name = os.path.join(current_dir, "localization_result_{}_anchors_{}.txt".format(6, i))
        i += 1
    np.savetxt(est_file_name, result)

    # gt_file_name = os.path.join(current_dir, "localization_result_{}_anchors_gt.txt".format(anchors_num))
    # if (os.path.exists(gt_file_name) == False):
    #     np.savetxt(gt_file_name, path_points)


def ekf_loop(path, data):
    estimation_result= []
    state_init = np.array([points[0,0], points[0,1], points[0,2], 0.0])
    print("state_init: {}".format(state_init))
    ekf_filter = heading_filter_tightly_couple_uwb(anchors, tags, state_init)
    for i in range(len(measurements_noise)):
        ekf_filter.predict()
        ekf_filter.roll = orient_noise[i,0]
        ekf_filter.pitch = orient_noise[i,1]
        ekf_filter.update(measurements_noise[i])
        # ekf_filter.roll = orient[i,0]
        # ekf_filter.pitch = orient[i,1]
        # ekf_filter.update(measurements[i])
        estimation_result.append(ekf_filter.state)
    # measurements_np = np.asarray(measurements).reshape(len(measurements), tags.shape[0] * anchors.shape[0])
    estimation_result_np = np.asarray(estimation_result).reshape(len(estimation_result), 4)
    # plot_result_style_1(estimation_result_np)
    save_result(path, estimation_result_np)


if __name__ == '__main__':

    # 解析命令行参数
    parser = argparse.ArgumentParser(description="Start the Optimization.")
    parser.add_argument('--experiment_path', type=str, default='None', help='anchor config file(default: None)')
    args = parser.parse_args()
    json_path = args.experiment_path
    # uwb1, uwb2 and uav imu & usv imu.
    if not os.path.exists(json_path):
        raise FileExistsError
    uwb1_data = load_json_file(os.path.join(json_path, "nodeframe3_tuav6_nlink_linktrack_nodeframe3.json"))
    uwb2_data = load_json_file(os.path.join(json_path, "nodeframe3_tuav62_nlink_linktrack_nodeframe3.json"))
    uav_data = load_json_file(os.path.join(json_path, "imu_tuav6_dji_osdk_ros_imu.json"))
    usv_data = load_json_file(os.path.join(json_path, "imu_usv_compass.json"))

    anchors = np.array([[-1.73, -2.4, 0.3],
                        [-1.73, -0.5, 1.6],
                        [-1.73, 0.5, 1.6],
                        [-1.73, 2.4, 0.3],
                        [1.73, 2.4, 0.3],
                        [1.73, -2.4, 0.3]])
    tags = np.array([[0.0, 0.0, 0.316],
                     [0.0, 0.0, -0.316]])
    
    synchronized_data = synchronize_data(uwb1_data, uwb2_data, uav_data, usv_data)
    ekf_loop(json_path, synchronized_data)
    
    